#include "YoloDet.h"


YoloDet::YoloDet()
{
}


YoloDet::~YoloDet()
{
}

int YoloDet::init(std::string model_path, std::string class_path, YoloParam param, bool is_cuda)
{
	if (mInited) 
	{
		std::cout << "alread init" << std::endl;
		return 0;
	}
	int nRet = -1;
	nRet = load_net(model_path, mNet, is_cuda);
	if (nRet) 
	{
		std::cout << "load_net err" << std::endl;
		return -1;
	}
	
	nRet = load_class_list(class_path, mClassNames);
	if (nRet) 
	{
		std::cout << "load_class_list err" << std::endl;
		return -1;
	}

	SCORE_THRESHOLD = param.score_threshold;		
	NMS_THRESHOLD =  param.nms_threshold;			
	CONFIDENCE_THRESHOLD = param.confidence_threshold;		
	return 0;
}

int  YoloDet::load_net(std::string str_path, cv::dnn::Net& net, bool is_cuda)
{
	if (str_path.empty())
	{
		std::cout << " net path is empty!" << std::endl;
		return -INFO_LOAD_NET_ERR;
	}
	
	net = cv::dnn::readNet(str_path);
	if (is_cuda)
	{
		net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
		net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16);
		std::cout << "Running on CUDA" << std::endl;
	}
	else
	{
		net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
		net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
		std::cout << "Running on CPU" << std::endl;
	}
	return INFO_OK;
}

int  YoloDet::load_class_list(std::string str_path, std::vector<std::string>&vec_class)
{
	if (str_path.empty())
	{
		std::cout << " class path is empty " << std::endl;
		return INFO_LOAD_CLASS_ERR;
	}

	vec_class.clear();
	std::ifstream ifs(str_path);
	std::string line;

	while (getline(ifs, line))
	{
		vec_class.push_back(line);
	}
	return INFO_OK;
}

cv::Mat YoloDet::format_yolov5(const cv::Mat& input)
{
	int col = input.cols;
	int row = input.rows;
	int _max = MAX(col, row);
	cv::Mat result = cv::Mat::zeros(_max, _max, CV_8UC3);
	input.copyTo(result(cv::Rect(0, 0, col, row)));
	return result;
}

int YoloDet::detect(cv::Mat& image, std::vector<YoloDetection>& output)
{
	cv::Mat blob;
    auto input_image = format_yolov5(image);

    cv::dnn::blobFromImage(input_image, blob, 1. / 255., \
                cv::Size(INPUT_WIDTH, INPUT_HEIGHT), cv::Scalar(), true, false);
	mNet.setInput(blob);

	std::vector<cv::Mat> outputs;

	mNet.forward(outputs, mNet.getUnconnectedOutLayersNames());

	float x_factor = input_image.cols / INPUT_WIDTH;
	float y_factor = input_image.rows / INPUT_HEIGHT;

	float* data = (float*)outputs[0].data;

		
	const int rows = 25200;			//网络最后一层结果的行数
	const int dimensions = 6;       //每个结果数据的长度        

	std::vector<int> class_ids;
	std::vector<float> confidences;
	std::vector<cv::Rect> boxes;
	for (int i = 0; i < rows; ++i)
	{
		float confidence = data[4];
		if (confidence >= CONFIDENCE_THRESHOLD)
		{
			float* classes_scores = data + 5;
            cv::Mat scores(1, (int)mClassNames.size(), CV_32FC1, classes_scores);
			cv::Point class_id;
			double max_class_score;
			minMaxLoc(scores, 0, &max_class_score, 0, &class_id);
			if (max_class_score > SCORE_THRESHOLD)
			{
				confidences.push_back(confidence);
				class_ids.push_back(class_id.x);

				float x = data[0];
				float y = data[1];
				float w = data[2];
				float h = data[3];

				int left = int((x - 0.5 * w) * x_factor);
				int top = int((y - 0.5 * h) * y_factor);
				int width = int(w * x_factor);
				int height = int(h * y_factor);
				boxes.push_back(cv::Rect(left, top, width, height));
			}
		}
		data += dimensions;
	}
	std::cout << "boxes.size() ="<<boxes.size() << std::endl;


	std::vector<int> nms_result;
	cv::dnn::NMSBoxes(boxes, confidences, SCORE_THRESHOLD, NMS_THRESHOLD, nms_result);
	for (int i = 0; i < nms_result.size(); i++)
	{
		int idx = nms_result[i];
		YoloDetection result;
		result.class_id = class_ids[idx];
		result.confidence = confidences[idx];
		result.box = boxes[idx];
		output.push_back(result);
	}

	std::cout << "output.size() =" << output.size() << std::endl;
	return INFO_OK;
}
